#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
数据生成器使用示例
演示如何使用data_gen.py进行数据处理
"""

import os
import numpy as np
from PIL import Image
from data_gen import data_flow, BaseSequence

def create_sample_data():
    """创建示例数据用于测试"""
    print("创建示例数据...")
    
    # 创建数据目录
    data_dir = "sample_data"
    os.makedirs(data_dir, exist_ok=True)
    
    # 创建示例图像和标签
    sample_images = [
        ("sample1.jpg", 0),  # 其他垃圾
        ("sample2.jpg", 6),  # 厨余垃圾
        ("sample3.jpg", 14), # 可回收物
        ("sample4.jpg", 37), # 有害垃圾
    ]
    
    for img_name, label in sample_images:
        # 创建随机图像
        img_array = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
        img = Image.fromarray(img_array)
        img_path = os.path.join(data_dir, img_name)
        img.save(img_path)
        
        # 创建标签文件
        label_path = os.path.join(data_dir, f"{img_name}.txt")
        with open(label_path, 'w', encoding='utf-8') as f:
            f.write(f"{img_name}, {label}")
    
    print(f"示例数据已创建在 {data_dir}/ 目录")
    return data_dir

def test_data_flow():
    """测试数据流生成器"""
    print("\n测试数据流生成器...")
    
    # 创建示例数据
    data_dir = create_sample_data()
    
    # 使用数据流生成器
    train_sequence, validation_sequence = data_flow(
        train_data_dir=data_dir,
        batch_size=2,
        num_classes=40,
        input_size=224
    )
    
    print(f"训练序列长度: {len(train_sequence)}")
    print(f"验证序列长度: {len(validation_sequence)}")
    
    # 获取一个批次的数据
    batch_x, batch_y = train_sequence[0]
    print(f"批次数据形状: {batch_x.shape}")
    print(f"批次标签形状: {batch_y.shape}")
    
    # 显示标签信息
    for i, label in enumerate(batch_y):
        pred_class = np.argmax(label)
        print(f"样本 {i}: 类别 {pred_class}")
    
    return train_sequence, validation_sequence

def test_custom_sequence():
    """测试自定义数据序列"""
    print("\n测试自定义数据序列...")
    
    # 准备数据
    img_paths = ["sample_data/sample1.jpg", "sample_data/sample2.jpg"]
    labels = np.array([[1, 0, 0, 0], [0, 1, 0, 0]])  # 2个类别，4个样本
    
    # 创建自定义序列
    custom_sequence = BaseSequence(
        img_paths=img_paths,
        labels=labels,
        batch_size=2,
        img_size=[224, 224]
    )
    
    print(f"自定义序列长度: {len(custom_sequence)}")
    
    # 获取数据
    batch_x, batch_y = custom_sequence[0]
    print(f"自定义批次数据形状: {batch_x.shape}")
    print(f"自定义批次标签形状: {batch_y.shape}")
    
    return custom_sequence

def test_preprocessing():
    """测试图像预处理"""
    print("\n测试图像预处理...")
    
    from models.resnet50 import preprocess_input
    
    # 创建测试图像
    test_img = np.random.randint(0, 255, (224, 224, 3), dtype=np.uint8)
    
    # 应用预处理
    processed_img = preprocess_input(test_img)
    
    print(f"原始图像范围: [{test_img.min()}, {test_img.max()}]")
    print(f"处理后图像范围: [{processed_img.min():.2f}, {processed_img.max():.2f}]")
    print(f"处理后图像形状: {processed_img.shape}")

def main():
    """主函数"""
    print("=" * 60)
    print("数据生成器使用示例")
    print("=" * 60)
    
    try:
        # 测试数据流
        train_seq, val_seq = test_data_flow()
        
        # 测试自定义序列
        custom_seq = test_custom_sequence()
        
        # 测试预处理
        test_preprocessing()
        
        print("\n" + "=" * 60)
        print("✅ 所有测试通过！")
        print("=" * 60)
        
        print("\n使用建议:")
        print("1. 确保数据格式正确 (图像文件 + 对应的.txt标签文件)")
        print("2. 调整batch_size以适应内存大小")
        print("3. 使用多进程数据加载提高训练效率")
        print("4. 可以通过修改BaseSequence类自定义预处理逻辑")
        
    except Exception as e:
        print(f"\n❌ 测试失败: {e}")
        print("请检查数据格式和依赖包是否正确安装")

if __name__ == "__main__":
    main() 